package com.xsherl.leetcode.solution;

import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Set;

public class TopKFrequent {

    /**
     * 给一非空的单词列表，返回前 k 个出现次数最多的单词。
     *
     * 返回的答案应该按单词出现频率由高到低排序。如果不同的单词有相同出现频率，按字母顺序排序。
     */
    public List<String> topKFrequent(String[] words, int k) {
        HashMap<String, Integer> wordCount = new HashMap<>();
        for (String word: words){
            int c = wordCount.getOrDefault(word, 0);
            wordCount.put(word, c + 1);
        }
        System.out.println(wordCount.get("hfhlfqzkcn"));
        System.out.println(wordCount.get("zddeno"));
        Set<String> wordSet = wordCount.keySet();
        String[] heap = new String[k];
        int i = 0;
        // 维持一个最小堆，保证堆里的元素是Top K里面的元素
        for (String word: wordSet){
            if (i < k){     //  构建最小堆
                heap[i] = word;
                swim(wordCount, heap, 0, i);
                i++;
            } else {
                // 如果发现比顶部还大的元素，就替换他
                if (greater(word, wordCount.get(word), heap[0], wordCount.get(heap[0]))){
                    heap[0] = word;
                    sink(wordCount, heap, 0, i - 1);
                }
            }
        }
        i--;
        while (i > 0){
            swap(heap, i, 0);
            i--;
            sink(wordCount, heap, 0, i);
        }
        return Arrays.asList(heap);
    }

    private boolean greater(String s1, int count1, String s2, int count2){
        return count1 > count2 || (count1 == count2 && s1.compareTo(s2) < 0);
    }

    private void sink(HashMap<String, Integer> wordCount, String[] heap, int begin, int end){
        int i = begin;
        while ((i * 2 + 1) <= end){
            int minIndex = i, left = i * 2 + 1, right = i * 2 + 2;
            if (greater(heap[minIndex],  wordCount.get(heap[minIndex]), heap[left], wordCount.get(heap[left]))){
                minIndex = left;
            }
            if (right <= end && greater(heap[minIndex],  wordCount.get(heap[minIndex]), heap[right], wordCount.get(heap[right]))){
                minIndex = right;
            }
            if (minIndex == i){
                break;
            }
            swap(heap, i, minIndex);
            i = minIndex;
        }
    }

    private void swim(HashMap<String, Integer> wordCount,String[] heap, int begin, int end){
        String s = heap[end];
        int c = wordCount.get(s);
        int i = end;
        while (((i - 1) / 2) >= begin && greater(heap[(i - 1) / 2], wordCount.get(heap[(i - 1) / 2]), s, c)){
            swap(heap, i, (i - 1) / 2);
            i = (i - 1) / 2;
        }
    }

    private void swap(String[] heap, int i, int j){
        String t = heap[i];
        heap[i] = heap[j];
        heap[j] = t;
    }
}
